# -*- coding:utf-8 -*-

# Copyright (C) 2020. Huawei Technologies Co., Ltd. All rights reserved.
# This program is free software; you can redistribute it and/or modify
# it under the terms of the MIT License.
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# MIT License for more details.

"""Conv module for backbone."""
import warnings
import torch.nn as nn
from .weight_init import kaiming_init, constant_init
from .norm import build_norm_layer

conv_cfg = {
    'Conv': nn.Conv2d,
}


def build_conv_layer(cfg, *args, **kwargs):
    """Build convolution layer.

    :param cfg: define how to build a conv layer
    :type: None or dict

    :return: created conv layer
    :rtype: nn.Module
    """
    if cfg is None:
        cfg_ = dict(type='Conv')
    else:
        # assert isinstance(cfg, dict) and 'type' in cfg
        cfg_ = cfg.copy()

    layer_type = cfg_.pop('type')
    if layer_type not in conv_cfg:
        raise KeyError('Unrecognized norm type {}'.format(layer_type))
    else:
        conv_layer = conv_cfg[layer_type]

    layer = conv_layer(*args, **kwargs, **cfg_)

    return layer


class ConvModule(nn.Module):
    """Create a conv block that contains conv/norm/activation layers.

    :param in_channels: Same as nn.Conv2d.
    :type: int
    :param out_channels: Same as nn.Conv2d.
    :type: int
    :param kernel_size: Same as nn.Conv2d.
    :type: int
    :param stride: Same as nn.Conv2d.
    :type: int or tuple[int]
    :param padding: Same as nn.Conv2d.
    :type: int or tuple[int]
    :param dilation: Same as nn.Conv2d.
    :type: int or tuple[int]
    :param groups: Same as nn.Conv2d.
    :type: int
    :param bias: If specified as `auto`, it will be decided by the norm_cfg.
                 Bias will be set as True if norm_cfg is None, otherwise False.
    :type: bool or str
    :param conv_cfg: Config dict for convolution layer.
    :type: dict
    :param norm_cfg: Config dict for normalization layer.
    :type: dict
    :param activation: Activation type, "ReLU" by default.
    :type: str or None
    :param inplace: Whether to use inplace mode for activation.
    :type: bool
    :param order: The order of conv/norm/activation layers. It is a sequence of "conv", "norm" and "act".
                  Examples are ("conv", "norm", "act") and ("act", "conv", "norm").
    :type: tuple[str]
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size,
                 stride=1,
                 padding=0,
                 dilation=1,
                 groups=1,
                 bias='auto',
                 conv_cfg=None,
                 norm_cfg=None,
                 activation='relu',
                 inplace=True,
                 order=('conv', 'norm', 'act')):
        super(ConvModule, self).__init__()
        # assert conv_cfg is None or isinstance(conv_cfg, dict)
        # assert norm_cfg is None or isinstance(norm_cfg, dict)
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.activation = activation
        self.inplace = inplace
        self.order = order
        # assert isinstance(self.order, tuple) and len(self.order) == 3
        # assert set(order) == set(['conv', 'norm', 'act'])

        self.with_norm = norm_cfg is not None
        self.with_activatation = activation is not None
        # if the conv layer is before a norm layer, bias is unnecessary.
        if bias == 'auto':
            bias = False if self.with_norm else True
        self.with_bias = bias

        if self.with_norm and self.with_bias:
            warnings.warn('ConvModule has norm and bias at the same time')

        # build convolution layer
        self.conv = build_conv_layer(
            conv_cfg,
            in_channels,
            out_channels,
            kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=groups,
            bias=bias)
        # export the attributes of self.conv to a higher level for convenience
        self.in_channels = self.conv.in_channels
        self.out_channels = self.conv.out_channels
        self.kernel_size = self.conv.kernel_size
        self.stride = self.conv.stride
        self.padding = self.conv.padding
        self.dilation = self.conv.dilation
        self.transposed = self.conv.transposed
        self.output_padding = self.conv.output_padding
        self.groups = self.conv.groups

        # build normalization layers
        if self.with_norm:
            # norm layer is after conv layer
            if order.index('norm') > order.index('conv'):
                norm_channels = out_channels
            else:
                norm_channels = in_channels
            self.norm_name, norm = build_norm_layer(norm_cfg, norm_channels)
            self.add_module(self.norm_name, norm)

        # build activation layer
        if self.with_activatation:
            # TODO: introduce `act_cfg` and supports more activation layers
            if self.activation not in ['relu']:
                raise ValueError('{} is currently not supported.'.format(
                    self.activation))
            if self.activation == 'relu':
                self.activate = nn.ReLU(inplace=inplace)

        # Use msra init by default
        self.init_weights()

    @property
    def norm(self):
        """Norm weight to init."""
        return getattr(self, self.norm_name)

    def init_weights(self):
        """Init weight."""
        nonlinearity = 'relu' if self.activation is None else self.activation
        kaiming_init(self.conv, nonlinearity=nonlinearity)
        if self.with_norm:
            constant_init(self.norm, 1, bias=0)

    def forward(self, x, activate=True, norm=True):
        """Calcation forward result."""
        for layer in self.order:
            if layer == 'conv':
                x = self.conv(x)
            elif layer == 'norm' and norm and self.with_norm:
                x = self.norm(x)
            elif layer == 'act' and activate and self.with_activatation:
                x = self.activate(x)
        return x
